#!/usr/bin/env python
import sys
import numpy as np
from glob import glob

t=float(sys.argv[2])
k=sys.argv[1]
convert = 0.01036427230133138/8.617333262145e-5/t

if __name__ == '__main__':


    head = 1000
    stride = 20

    newdata = []
    # ff = open("da", "w")
    allxi = []
    allcolvar = []
    allpe = []
    print("all folders", f"../../umb_{k}_*/*.npz", glob(f"../../umb_{k}_*/*.npz"))
    for filename in glob(f"../../umb_{k}*/*.npz"):
        data = np.load(filename)
        xi = data['xi']
        e = data['alle'][:, 1]
        if xi.shape[0]>0 and np.max(e)<5:

            intc = data['colvar']

            d = {}
            d['xi'] = xi[head::stride]
            d['colvar'] = intc[head::stride, 0:2]
            d['x0'] = data['center']
            d['k'] = data['k']
            if 'k0' in data:
                d['k0'] = data['k0']
            else:
                d['k0'] = 0.
            print(filename, len(d['xi']))

            allxi += [xi[::stride]]
            allcolvar += [intc[::stride]]
            allpe += [e[::stride]]

            # phi = idata['xi']
            # print(np.average(phi), d['x0'], d['k'], file=ff)
            newdata += [d]

    np.savez(f"{k}.npz", alldata=newdata)
    allxi = np.hstack(allxi)
    allpe = np.hstack(allpe)
    allcolvar = np.vstack(allcolvar)
    np.savez(f"{k}_sum.npz", xi=allxi, pe=allpe, colvar=allcolvar)

    alldata = newdata
    f = open(f"{k}.info", "w")
    for i, idata in enumerate(alldata):
        phi = idata['xi']
        print(f"# state {i:d}", file=f)
        biasE = []
        for j, jdata in enumerate(alldata):
            biasE += [0.5*jdata['k']*(phi-jdata['x0'])**2*convert]
        biasE = np.vstack(biasE)
        biasE = biasE.T
        np.savetxt(f, biasE)
